import json
import os

os.environ["MUJOCO_GL"] = "egl"
os.environ["ROS_PACKAGE_PATH"] = "piper_description"

import mediapy as media
import numpy as np
import pinocchio as pin
from gymnasium import spaces
from scipy.spatial.transform import Rotation
from tqdm import tqdm

from sim.piper_simulator import MujocoPiperSim

# def random_rpy_round_anchor(anchor_rpy: np.ndarray = np.array([0, 2, 0]), std: float = 1):
#     """
#     Gaussian sampling of the rpy with an anchor
#     The anchor_rpy is the initial rpy of the effector, which is pointing forward & slightly down
#     """
#     return np.random.normal(anchor_rpy, std)


def random_quaternion_round_anchor(anchor_quat: np.ndarray = np.array([0.4, 0, 0.8, 0])):
    """
    uniform sampling of the quaternion with an anchor
    Shoemake method
    """
    u1, u2, u3 = np.random.uniform(0, 1, 3)

    # anchor the quaternion: effector is pointing forward & slightly down
    # cos(theta) = 0.4, sin(theta) = 0.8
    # the smaller the w, the larger the angle
    anchor_quat = anchor_quat

    # compute the quaternion components
    w = np.sqrt(1 - u1) * np.sin(2 * np.pi * u2)
    x = np.sqrt(1 - u1) * np.cos(2 * np.pi * u2)
    y = np.sqrt(u1) * np.sin(2 * np.pi * u3)
    z = np.sqrt(u1) * np.cos(2 * np.pi * u3)

    # t = 0.5: sampling the "forward sphere"
    return slerp(anchor_quat, np.array([w, x, y, z]), 0.5)


def slerp(q1, q2, t):
    """Spherical linear interpolation between two quaternions"""

    if np.dot(q1, q2) < 0:
        q2 = -q2

    # Calculate angle between quaternions
    cos_theta = np.dot(q1, q2)
    cos_theta = np.clip(cos_theta, -1.0, 1.0)
    theta = np.arccos(cos_theta)

    if theta < 1e-6:
        return q1

    # SLERP formula
    sin_theta = np.sin(theta)
    w1 = np.sin((1 - t) * theta) / sin_theta
    w2 = np.sin(t * theta) / sin_theta

    return w1 * q1 + w2 * q2


class CollisionDetector:
    def __init__(self):
        # just for collision check
        urdf = "piper_description/piper_description.urdf"
        self.pin_robot = pin.RobotWrapper.BuildFromURDF(urdf)
        self.geom_model = pin.buildGeomFromUrdf(self.pin_robot.model, urdf, pin.GeometryType.COLLISION)
        for i in range(4, 9):
            for j in range(0, 3):
                self.geom_model.addCollisionPair(pin.CollisionPair(i, j))
        self.geometry_data = pin.GeometryData(self.geom_model)

    def __call__(self, joint):
        """
        check whether the arm collide with itself
        """
        pin.forwardKinematics(self.pin_robot.model, self.pin_robot.data, joint)
        pin.updateGeometryPlacements(self.pin_robot.model, self.pin_robot.data, self.geom_model, self.geometry_data)
        collision = pin.computeCollisions(self.geom_model, self.geometry_data, False)

        # print("collision:", collision)
        return collision


class TrajectorySampler:
    """
    stepsize: the step size of the trajectory in xyz space
    radius: the constraint of the orientation residual
    reg: the regularization weight for smoothing the trajectory
    """

    def __init__(self, stepsize=0.1, radius=0.04, reg=0.01, **kwargs):
        self.ik = MujocoPiperSim(radius=radius, reg=reg, **kwargs)
        self.stepsize = stepsize  # in meters

        # the sampling spaces: using full space?
        max_pos_m = np.array([0.7, 0.7, 0.7])
        min_pos_m = np.array([-0.1, -0.7, -0.1])
        # max_pos_m = np.array([0.6, 0.6, 0.6])
        # min_pos_m = np.array([-0.01, -0.6, -0.01])
        self.xspace = spaces.Box(low=min_pos_m, high=max_pos_m, shape=(3,))
        self.jspace = spaces.Box(low=self.ik.bounds[0][:6], high=self.ik.bounds[1][:6], shape=(6,))

        self.collision_detector = CollisionDetector()

    def get_line_xyz_trajectory(self):
        tar_pos, tar_quat = self.xspace.sample(), random_quaternion_round_anchor()
        pos, quat = self.ik.get_effector_xquat()[:3], self.ik.get_effector_xquat()[3:]
        n_frame = int(np.linalg.norm(pos - tar_pos) / 0.01) + 1
        t = 0
        while t <= n_frame:
            pos_ = pos + (tar_pos - pos) * t / n_frame
            quat_ = slerp(quat, tar_quat, t / n_frame)
            t += 1
            yield pos_, quat_

    def get_joint_trajectory(self):
        tar_joint = self.jspace.sample()
        joint = self.ik.j_home[:6].copy()
        n_frame = int(np.linalg.norm(joint - tar_joint) / 0.05) + 1
        t = 0
        while t <= n_frame:
            joint_ = joint + (tar_joint - joint) * t / n_frame
            t += 1
            yield joint_

    def sample_trajectory(self, need_frames=False):
        """
        Sample a trajectory of the end effector using position and orientation.
        The intermediate positions are sampled by linear interpolation.
        The quaternion is sampled by spherical linear interpolation， which is more intuitive

        The actions uses xyz and rpy

        No storing the joints rightnow
        """
        tar_x, tar_quat = self.xspace.sample(), random_quaternion_round_anchor()

        if np.random.rand() < 0.01:
            self.ik.go_zero()

        cur_x, cur_quat = self.ik.get_effector_xquat()[:3], self.ik.get_effector_xquat()[3:]

        n_frame = int(np.linalg.norm(cur_x - tar_x) / self.stepsize) + 1  # at least 1 frame = one-step

        joints = []
        xeuler = []
        frames = []
        t = 0  # there are n_frame + 1 joints: including the initial position
        j = self.ik.data.qpos.copy()  # initial guess=current state
        while t <= n_frame:
            pos = cur_x + (tar_x - cur_x) * t / n_frame
            quat = slerp(cur_quat, tar_quat, t / n_frame)
            j_prev = j.copy()  # continuity regularization
            jo = self.ik.solve_ik_xquat(pos=pos, quat=quat, j0=j, x_prev=j_prev)

            # record the joint angles and the xquat
            joints.append(jo)
            self.ik.step()
            # using exact xquat instead of the target(imagined) xquat
            xeuler.append(self.ik.get_effector_xeuler())

            if need_frames:
                frames.append(self.ik.render_frame())
            t += 1

        # check collision
        if any(self.collision_detector(j) for j in joints):
            print("Collision detected")
            return None

        return {
            "xeuler": [_.tolist() for _ in xeuler],
            "frames": frames,
        }

    def create_data(self, num_traj: int):

        data = {"xeuler": [], "lengths": []}
        num_data = 0
        filename = f"data/json_data/sim_data_1Mfull_{self.stepsize}.json"
        progress_bar = tqdm(desc="Generate traj.", unit="it", total=num_traj)
        while num_data < num_traj:
            traj = self.sample_trajectory()
            if traj:
                data["xeuler"] += traj["xeuler"]
                data["lengths"].append(len(traj["xeuler"]))
                num_data += 1
                progress_bar.update(1)

        progress_bar.close()
        with open(filename, "w") as f:
            json.dump(data, f)

        print(f"Data saved to {filename}")


if __name__ == "__main__":
    """
    sampling:
    PYTHONPATH=/home/jj/knowin-robort python sim/sampler.py
    """

    # sampler = TrajectorySampler(
    #     stepsize=0.1,
    #     radius=0.04,
    #     reg=0.1,
    # )
    # sampler.create_data(num_traj=1000_000)

    def test_sampler():
        # upper: the joint (true) plot
        # lower: the xyz and rpy (ik-solved) plot
        # expected the same behaviors: the ik solution is always valid
        self = MujocoPiperSim(height=240, width=480)

        sampler = TrajectorySampler(stepsize=0.1, radius=0.04, reg=0.1, height=240, width=480)
        frames = None

        i = 0
        while i < 10:
            traj = sampler.sample_trajectory(need_frames=True)
            if traj is not None:
                j_frames = np.stack(traj["frames"])
                # frames += traj["frames"]
                d_frames = []
                for xe in traj["xeuler"]:
                    self.move_to(xe[:3], xe[3:])
                    d_frames.append(self.render_frame())
                d_frames = np.stack(d_frames)
                if frames is None:
                    frames = np.concat([j_frames, d_frames], axis=1)
                else:
                    frames = np.concat([frames, np.concat([j_frames, d_frames], axis=1)], axis=0)
                i += 1

        media.show_video(frames, loop=False)

    def test_ik_line_sampling():
        sampler = TrajectorySampler()
        frames = []
        j = sampler.ik.j_home.copy()
        for i in range(20):
            for pos, quat in sampler.get_line_xyz_trajectory():
                j_prev = sampler.ik.j_home.copy()  # update the previous joint angles
                j = sampler.ik.solve_ik_xquat(pos=pos, quat=quat, j0=j, x_prev=j_prev)
                sampler.ik.step()
                frames.append(sampler.ik.render_frame())
        media.show_video(frames, loop=False)

    def test_ik_joint_sampling():
        # joint sampling: trajectory is dangerous: the arm usually swings fast.
        sampler = TrajectorySampler()
        frames = []
        for i in range(20):
            for joint in sampler.get_joint_trajectory():
                sampler.ik.data.qpos = np.concat((joint, np.zeros(2)))
                sampler.ik.step()
                frames.append(sampler.ik.render_frame())

        media.show_video(frames, loop=False)

    def test_mj_simulator():
        """
        Instructions for the Piper arm:

        1. solve the joint angles:

        j = self.solve_ik_xyz_rpy(pos=pos, rpy=rpy, j0=j0, x_prev=j0.copy())

        2. transform the joint angles to the machine angles:

            machine_joints = (np.rad2deg(joints) * 1000).astype(int)
            arm.ModeCtrl(0x01, 0x01, self.__speed, 0x00)
            arm.JointCtrl(*machine_joints)

        """
        self = MujocoPiperSim()
        xyz = np.array([0.34803, 0.27306, 0.2])  # in meters
        rpy = np.array([0, 1.7, 0])  # in radians
        j0 = self.j_home.copy()  # initial guess in radians
        j = self.solve_ik_xyz_rpy(xyz=xyz, rpy=rpy, j0=j0, x_prev=j0.copy())

        # check
        self.step(j)
        media.show_image(self.render_frame())
